from stable_baselines3.common.callbacks import BaseCallback

class InfoAveragingCallback(BaseCallback):
    def __init__(self, eval_env, eval_freq=5000, n_eval_episodes=5, verbose=0):
        super().__init__(verbose)
        self.eval_env = eval_env
        self.eval_freq = eval_freq
        self.n_eval_episodes = n_eval_episodes

    def _on_step(self) -> bool:
        if self.num_timesteps % self.eval_freq != 0:
            return True
        import numpy as np
        stats = {"chi2": [], "L": [], "cond": [], "N": []}
        obs, _ = self.eval_env.reset()
        for _ in range(self.n_eval_episodes):
            done = False
            while not done:
                action, _ = self.model.predict(obs, deterministic=True)
                obs, r, done, trunc, info = self.eval_env.step(action)
            # 每个episode收尾时记一次
            for k in stats: stats[k].append(info[k] if k in info else None)
            obs, _ = self.eval_env.reset()
        for k,v in stats.items():
            vals = [x for x in v if x is not None]
            if vals:
                self.logger.record(f"biz/{k}", float(np.mean(vals)))  # ★ biz标签
        return True
